import json

from pyflink.common import SimpleStringSchema, WatermarkStrategy, Types
from pyflink.datastream import StreamExecutionEnvironment
from pyflink.datastream.connectors import DeliveryGuarantee
from pyflink.datastream.connectors.kafka import KafkaSource, KafkaOffsetsInitializer, KafkaSink, \
    KafkaRecordSerializationSchema

# 1、创建环境
env = StreamExecutionEnvironment.get_execution_environment()

# 并行度设置
# 设置方式
# 1、在代码中设置，优先级最高
# 2、提交任务是设置 -p
# 3、配置文件中有默认值，默认是1，parallelism.default: 1

# 设置原则，基于数据量设置（每秒的数据量）
# 1、非聚合计算，flink单个并行度的吞吐量大概10万/s
# 2、聚合计算，flink单个并行度吞吐量大概1万/s
# 3、也可以参考kafka的分区数来设置并行度

# flink资源
# flink默认一个并行度对应一个slot
# 设置多少并行度就需要多少资源

# env.set_parallelism(1)

# 2、读取kafka中的数据
source = KafkaSource.builder() \
    .set_bootstrap_servers("master:9092") \
    .set_topics("students_partition_age") \
    .set_group_id("my-group") \
    .set_starting_offsets(KafkaOffsetsInitializer.earliest()) \
    .set_value_only_deserializer(SimpleStringSchema()) \
    .build()

# 基于kafka source构建流（无界流）
students_ds = env.from_source(source, WatermarkStrategy.no_watermarks(), "Kafka Source")


# 解析数据
def map_fun(line):
    clazz = json.loads(line)["clazz"]
    return clazz, 1


clazz_ds = students_ds.map(map_fun)

# 统计班级的人数
clazz_num_ds = clazz_ds.key_by(lambda kv: kv[0]).sum(1)

# 整理格式
result_ds = clazz_num_ds.map(lambda kv: f"{kv[0]},{kv[1]}", output_type=Types.STRING())

# 2、将计算结果保存到kafka中
# set_bootstrap_servers: broker列表
# set_record_serializer: 指定topic和数据格式
# set_delivery_guarantee:数据处理的语义，AT_LEAST_ONCE：至少一次，EXACTLY_ONCE:唯一一次
sink = KafkaSink.builder() \
    .set_bootstrap_servers("master:9092") \
    .set_record_serializer(
    KafkaRecordSerializationSchema.builder()
    .set_topic("clazz_num")
    .set_value_serialization_schema(SimpleStringSchema())
    .build()) \
    .set_delivery_guarantee(DeliveryGuarantee.AT_LEAST_ONCE) \
    .build()

result_ds.sink_to(sink)

env.execute()